import sys
import os
from PyQt5 import QtGui
from PyQt5.QtWidgets import QApplication, QMainWindow, QLabel, QPushButton
from PyQt5.QtGui import QPixmap, QImage, QPainter, QImageReader
from PyQt5.QtCore import Qt, QRect

import detector

class ImageSlicer(QMainWindow):
    def __init__(self):
        try:
            super().__init__()

            self.image_folders = [
                "../fabric-defect/mapi_Standardize",
                "../fabric-defect/shousun+jinjing_Standardize",
                "../fabric-defect/youwu+quewei_Standardize"
            ]
            self.image_list = self.get_image_list()
            self.current_index = 0
            self.image_path = self.image_list[0]
            self.image = QImage(self.image_path)
            self.selected_area = QRect()
            self.saved_slices_count = 0  # 保存的切片数量
            self.initUI()
        except Exception as e:
            print(f"Error in __init__: {e}")

    def initUI(self):
        self.setWindowTitle("Image Slicer")

        self.label = QLabel(self)
        self.label.setGeometry(10, 50, self.image.width(), self.image.height())
        # self.label.setAlignment(Qt.AlignTop)

        self.label_path = QLabel(self)
        self.label_path.setGeometry(10, 10, 1500, 50)
        self.label_path.setAlignment(Qt.AlignTop)

        self.selection_pen = QtGui.QPen(Qt.green, 3)  # 颜色，线宽

        self.confirm_button = QPushButton("Confirm Selection", self)
        self.confirm_button.setGeometry(1000, 50 + self.image.height() + 20, 500, 200)
        self.confirm_button.clicked.connect(self.confirmSelection)

        self.prev_button = QPushButton("Previous", self)
        self.prev_button.setGeometry(300, 50 + self.image.height() + 20, 500, 200)
        self.prev_button.clicked.connect(self.show_previous_image)

        self.next_button = QPushButton("Next", self)
        self.next_button.setGeometry(1600, 50 + self.image.height() + 20, 500, 200)
        self.next_button.clicked.connect(self.show_next_image)

        self.setMouseTracking(True) # 确保启用鼠标跟踪
        
        self.update_image()

    def mousePressEvent(self, event):
        if event.button() == Qt.LeftButton:
            self.selected_area.setTopLeft(event.pos())
            # 清除保存文件名的显示
            self.label_path.clear()

    def mouseMoveEvent(self, event):
        if event.buttons() & Qt.LeftButton:
            self.selected_area.setBottomRight(event.pos())
            self.updateSelectedArea()

    def mouseReleaseEvent(self, event):
        if event.button() == Qt.LeftButton:
            self.selected_area.setBottomRight(event.pos())
            self.updateSelectedArea()

    def updateSelectedArea(self):
        # 创建图像副本
        temp_image = self.image.copy()

        # 在副本上绘制选区
        painter = QPainter(temp_image)
        # 设置笔刷属性
        painter.setPen(self.selection_pen)
        painter.drawRect(self.selected_area)
        # 在 QLabel 中显示更新后的图像
        self.label.setPixmap(QPixmap.fromImage(temp_image))

        # 确保了在每次绘制完成后正确释放 QPainter 对象
        painter.end()

    def confirmSelection(self):
        # 获取 self.selected_area 的坐标信息, 实现保存选区的逻辑
        top = self.selected_area.topLeft().y()
        left = self.selected_area.topLeft().x()
        bottom = self.selected_area.bottomRight().y()
        right = self.selected_area.bottomRight().x()

        img = detector.cv_read(self.image_path)
        img_crop = detector.get_roi(img, [left,top], [right,bottom])

        new_slice_index = self.saved_slices_count
        new_full_path = detector.gen_new_path(self.image_path,"crop_normal",f"crop_normal_{new_slice_index:04d}")

        # 获取目标文件夹中已经保存的切片数量
        saved_slices_folder = os.path.join(os.path.dirname(new_full_path))
        existing_slices = [file for file in os.listdir(saved_slices_folder) if file.startswith("crop_normal")]
        existing_slices_count = len(existing_slices)
        # 为新的切片自动设定一个编号
        new_slice_index = existing_slices_count + 1
        new_full_path = detector.gen_new_path(self.image_path,"crop_normal",f"crop_normal_{new_slice_index:04d}")
        detector.save_img(new_full_path, img_crop)

        # 更新保存的切片数量
        self.saved_slices_count = new_slice_index
        
        self.label_path.setText(f"Saved Image: {new_full_path}")

    def get_image_list(self):
        image_list = []
        supported_formats = [format.data().decode() for format in QImageReader.supportedImageFormats()]
        for i in range(len(self.image_folders)):
            for root, dirs, files in os.walk(self.image_folders[i]):
                for file in files:
                    if file.lower().endswith(tuple(supported_formats)):
                        image_list.append(os.path.join(root, file))
        return image_list

    def update_image(self):
        if 0 <= self.current_index < len(self.image_list):
            self.image_path = self.image_list[self.current_index]
            self.setWindowTitle(f"Image Slicer - {self.image_path}")
            self.image = QImage(self.image_path)
            self.resize(round(self.image.width()*1.2), round(self.image.height()*1.15))
            self.image = self.image.convertToFormat(QImage.Format_RGB32)
            self.label.setPixmap(QPixmap.fromImage(self.image))

            # 清除保存文件名的显示
            self.label_path.clear()

    def show_previous_image(self):
        self.current_index -= 1
        if self.current_index < 0:
            self.current_index = len(self.image_list) - 1
        self.update_image()

    def show_next_image(self):
        self.current_index += 1
        if self.current_index >= len(self.image_list):
            self.current_index = 0
        self.update_image()

if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = ImageSlicer()
    window.show()
    sys.exit(app.exec_())